import os
import shutil

data_root = "../mnist_jpg"
def organize_images_by_prefix(source_folder, target_folder):
    # 创建目标文件夹（如果不存在的话）
    train_folder = os.path.join(target_folder, 'train')
    test_folder = os.path.join(target_folder, 'test')

    os.makedirs(train_folder, exist_ok=True)
    os.makedirs(test_folder, exist_ok=True)

    # 遍历源文件夹中的所有文件
    for filename in os.listdir(source_folder):
        file_path = os.path.join(source_folder, filename)

        if os.path.isfile(file_path):
            # 根据文件名前缀决定目标文件夹
            if filename.startswith('training_'):
                shutil.move(file_path, os.path.join(train_folder, filename))
            elif filename.startswith('test_'):
                shutil.move(file_path, os.path.join(test_folder, filename))
            else:
                print(f"File {filename} does not match 'train_' or 'test_' prefix.")

source_folder = 'path_to_your_source_folder'  # 替换为你的源文件夹路径
target_folder = 'path_to_your_target_folder'  # 替换为你的目标文件夹路径

organize_images_by_prefix(data_root, data_root)
